import json
import argparse
import sys
sys.path.append("../..")
sys.path.append("../../third_party/Matcha-TTS")
from websocket_util import WebSocketServer
from typing import Optional, Literal, Annotated
from pydantic import BaseModel, Field, ValidationError
from async_cosyvoice.async_cosyvoice import AsyncCosyVoice2
from websockets.legacy.server import WebSocketServerProtocol

class SpeechRequest(BaseModel):
    """语音合成请求参数"""
    input: str = Field(
        ...,
        max_length=4096,
        examples=["你好，欢迎使用语音合成服务！"], description="需要转换为语音的文本内容"
    )
    voice: Optional[str] = Field(
        "001",
        examples=[
            "001",
            "speech:voice-name:xxx:xxx",
        ],
        description="音色选择"
    )
    response_format: Optional[Literal["mp3", "wav"]] = Field(
        "mp3",
        examples=["mp3", "wav"],
        description="输出音频格式"
    )
    sample_rate: Optional[int] = Field(
        24000,
        description="采样率，目前不支持设置，默认为返回 24000 Hz音频数据"
    )
    stream: Optional[bool] = Field(
        False,
        description="暂不支持流式返回。"
    )
    speed: Annotated[Optional[float], Field(strict=True, ge=0.25, le=4.0)] = Field(
        1.0,
        description="语速控制[0.25-4.0]"
    )

class WSResult(BaseModel):
    """Websocket返回结构"""
    status: Optional[int] = Field(
        0,
        description="成功"
    )
    response: Optional[list] = Field(
        [],
        description="返回的语音数据"
    )
    message: Optional[str] = Field(
        "success",
        description="对返回结果的描述"
    )

class WSServer(WebSocketServer):
    def __init__(self, args: argparse.Namespace):
        self.cosyvoice = AsyncCosyVoice2(args.model_dir, load_jit=args.load_jit)
        super().__init__()

    # 消息处理函数
    async def receive_msg_handler(self, websocket: WebSocketServerProtocol):
        async for message in websocket:
            ws_error_result: str = WSResult(
                status=-1,
                message="请发送正确的Json数据,必填字段【input】为需要生成语音的文字。"
            ).model_dump_json()
            try:
                msg: dict = json.loads(message)
                speed_request: SpeechRequest = SpeechRequest(**msg)
            except json.decoder.JSONDecodeError:
                await websocket.send(ws_error_result)
                continue
            except ValidationError:
                await websocket.send(ws_error_result)
                continue
            tts_text: str = speed_request.input
            spk_id: str = speed_request.voice
            end_of_prompt_index: int = tts_text.find("<|endofprompt|>")
            if end_of_prompt_index != -1:
                instruct_text: str = tts_text[:end_of_prompt_index + len("<|endofprompt|>")]
                tts_text: str = tts_text[end_of_prompt_index + len("<|endofprompt|>"):]
                try:
                    async for chunk in self.cosyvoice.inference_instruct2_by_spk_id(
                            tts_text,
                            instruct_text,
                            spk_id,
                            stream=speed_request.stream,
                            speed=speed_request.speed,
                            text_frontend=True
                    ):
                        await websocket.send(WSResult(response=chunk['tts_speech'].tolist()).model_dump_json())
                    await websocket.send(WSResult(status=-2).model_dump_json())
                except Exception as e:
                    await websocket.send(WSResult(status=-1, message=f"识别报错：{e}").model_dump_json())
            else:
                try:
                    async for chunk in self.cosyvoice.inference_zero_shot_by_spk_id(
                            tts_text,
                            spk_id,
                            stream=speed_request.stream,
                            speed=speed_request.speed,
                            text_frontend=True
                    ):
                        await websocket.send(WSResult(response=chunk['tts_speech'].tolist()).model_dump_json())
                    await websocket.send(WSResult(status=-2).model_dump_json())
                except Exception as e:
                    await websocket.send(WSResult(status=-1, message=f"识别报错：{e}").model_dump_json())


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', type=str,
                        default='../../pretrained_models/CosyVoice2-0.5B',
                        help='local path or modelscope repo id')
    parser.add_argument('--load_jit', action='store_true', help='load jit model', default=True)
    args = parser.parse_args()
    WSServer(args)